import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from typing import TypedDict, Annotated, List, Union
from langgraph.graph import StateGraph, END
import os

# --- 0. Matplotlib 设置 ---
plt.rcParams['font.sans-serif'] = ['SimHei'] # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False   # 解决负号显示问题

# --- 1. 数据处理和模型函数 (与之前类似，稍作调整以适应LangGraph) ---
def load_data_from_path(file_path: str) -> pd.DataFrame | None:
    try:
        df = pd.read_csv(file_path)
        df.rename(columns={'date': 'ds', 'exchange_rate': 'y'}, inplace=True)
        df['ds'] = pd.to_datetime(df['ds'])
        df.sort_values('ds', inplace=True)
        df['y'] = pd.to_numeric(df['y'], errors='coerce')
        df.dropna(subset=['y'], inplace=True)
        print(f"数据加载成功，共 {len(df)} 条记录。")
        return df
    except FileNotFoundError:
        print(f"错误：数据文件未找到 - {file_path}")
        return None
    except Exception as e:
        print(f"加载数据时发生错误: {e}")
        return None

def create_features_for_model(df: pd.DataFrame, lag: int = 1) -> pd.DataFrame | None:
    if df is None or df.empty:
        return None
    df_copy = df.copy()
    df_copy[f'y_lag_{lag}'] = df_copy['y'].shift(lag)
    df_copy.dropna(inplace=True)
    return df_copy

def train_linear_regression_model(df_features: pd.DataFrame) -> LinearRegression | None:
    if df_features is None or df_features.empty or 'y' not in df_features.columns or not any(col.startswith('y_lag_') for col in df_features.columns):
        print("错误：特征数据不完整，无法训练模型。")
        return None
    X = df_features[[col for col in df_features.columns if col.startswith('y_lag_')]]
    y = df_features['y']
    model = LinearRegression()
    model.fit(X, y)
    print("线性回归模型训练完成。")
    return model

def predict_future_values(model: LinearRegression, last_known_value: float, n_days: int = 7, lag: int = 1) -> List[float]:
    if model is None:
        return []
    predictions = []
    current_input = last_known_value
    for _ in range(n_days):
        feature_vector = np.array([[current_input]]) # Assumes lag=1
        predicted_value = model.predict(feature_vector)[0]
        predictions.append(predicted_value)
        current_input = predicted_value
    return predictions

def plot_forecast(historical_df: pd.DataFrame, future_dates: List[pd.Timestamp], future_predictions: List[float], output_filename: str = 'langgraph_lr_forecast.png') -> str:
    plt.figure(figsize=(15, 7))
     # 调整历史数据显示范围，例如显示最后90天的数据
    if len(historical_df) > 90:
        display_historical_df = historical_df.iloc[-90:]
    else:
        display_historical_df = historical_df
    plt.plot(historical_df['ds'], historical_df['y'], label='历史汇率(最近90天)', color='blue')
    plt.plot(future_dates, future_predictions, label=f'预测汇率(未来{len(future_predictions)}天)', color='red', linestyle='--')
    plt.title(f'汇率预测 - LangGraph与线性回归(未来{len(future_predictions)}天)')
    plt.xlabel('日期')
    plt.ylabel('汇率')
    plt.legend()
    plt.grid(True)
    # 自动调整X轴的日期显示范围，以包含所有数据点
    all_dates = pd.concat([display_historical_df['ds'], pd.Series(future_dates)])
    if not all_dates.empty:
       plt.xlim(all_dates.min(), all_dates.max())
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
     # 根据总时间跨度动态调整日期间隔
    total_days_span = (all_dates.max() - all_dates.min()).days if not all_dates.empty else 30
    if total_days_span > 180:
        plt.gca().xaxis.set_major_locator(mdates.MonthLocator(interval=2))
    elif total_days_span > 60:
        plt.gca().xaxis.set_major_locator(mdates.MonthLocator(interval=1))
    else:
        plt.gca().xaxis.set_major_locator(mdates.DayLocator(interval=max(1, total_days_span // 10))) # 保证大约10个刻度
    plt.gcf().autofmt_xdate()
    try:
        plt.savefig(output_filename)
        print(f"图表已保存为: {output_filename}")
        return output_filename
    except Exception as e:
        print(f"保存图表时出错: {e}")
        return ""
    finally:
        plt.close() # 关闭图像，避免在循环中重复显示

# --- 2. 定义 LangGraph 状态 ---
class AgentState(TypedDict):
    data_file_path: str
    raw_data: pd.DataFrame | None
    feature_data: pd.DataFrame | None
    model: LinearRegression | None
    user_query: str | None
    days_to_predict: int
    predictions: List[float] | None
    prediction_dates: List[pd.Timestamp] | None
    plot_path: str | None
    next_action: str # 用于控制流程
    error_message: str | None

# --- 3. 定义 LangGraph 节点 ---
def initialize_agent(state: AgentState) -> AgentState:
    print("代理初始化...")
    state['data_file_path'] = r'd:/13langgraph/旧的/usd_cny.csv' # 固定文件路径
    state['days_to_predict'] = 30 # 默认预测7天
    state['next_action'] = 'load_data'
    state['error_message'] = None
    return state


def load_data_node(state: AgentState) -> AgentState:
    print("节点：加载数据...")
    df = load_data_from_path(state['data_file_path'])
    if df is not None:
        state['raw_data'] = df
        state['next_action'] = 'train_model'
        state['error_message'] = None
    else:
        state['error_message'] = "数据加载失败。请检查文件路径和内容。"
        state['next_action'] = 'handle_error' # 或者 'end_conversation'
    return state

def train_model_node(state: AgentState) -> AgentState:
    print("节点：训练模型...")
    if state['raw_data'] is None:
        state['error_message'] = "没有加载数据，无法训练模型。"
        state['next_action'] = 'handle_error'
        return state
    
    features = create_features_for_model(state['raw_data'], lag=1)
    state['feature_data'] = features
    model = train_linear_regression_model(features)
    if model is not None:
        state['model'] = model
        state['next_action'] = 'get_user_query'
        state['error_message'] = None
    else:
        state['error_message'] = "模型训练失败。"
        state['next_action'] = 'handle_error'
    return state

def get_user_query_node(state: AgentState) -> AgentState:
    print("\n请输入您的请求 (例如：'预测未来7天汇率', '预测15天', '退出'):")
    query = input("> ").strip().lower()
    state['user_query'] = query
    # 简单的查询解析
    if "退出" in query or "exit" in query:
        state['next_action'] = 'end_conversation'
    elif "预测" in query:
        try:
            # 尝试从查询中提取天数
            num_str = "".join(filter(str.isdigit, query))
            if num_str:
                state['days_to_predict'] = int(num_str)
                print(f"将预测未来 {state['days_to_predict']} 天的汇率。")
            else:
                state['days_to_predict'] = 7 # 默认值
                print(f"未指定明确天数，将预测未来 {state['days_to_predict']} 天的汇率。")
            state['next_action'] = 'predict_exchange_rate'
        except ValueError:
            print("无法解析预测天数，请使用如 '预测10天' 的格式。")
            state['next_action'] = 'get_user_query' # 重新提问
    else:
        print("无法理解您的请求，请重新输入。")
        state['next_action'] = 'get_user_query'
    return state

def predict_exchange_rate_node(state: AgentState) -> AgentState:
    print("节点：进行预测...")
    if state['model'] is None or state['raw_data'] is None or state['raw_data'].empty:
        state['error_message'] = "模型未训练或数据未加载，无法预测。"
        state['next_action'] = 'handle_error'
        return state

    last_known_value = state['raw_data']['y'].iloc[-1]
    last_known_date = state['raw_data']['ds'].iloc[-1]
    n_days = state.get('days_to_predict', 7)

    predictions = predict_future_values(state['model'], last_known_value, n_days=n_days, lag=1)
    if predictions:
        state['predictions'] = predictions
        state['prediction_dates'] = [last_known_date + pd.Timedelta(days=i) for i in range(1, n_days + 1)]
        print("--- 预测结果 ---")
        for date, val in zip(state['prediction_dates'], state['predictions']):
            print(f"{date.strftime('%Y-%m-%d')}: {val:.4f}")
        state['next_action'] = 'generate_plot'
        state['error_message'] = None
    else:
        state['error_message'] = "预测失败。"
        state['next_action'] = 'handle_error'
    return state

def generate_plot_node(state: AgentState) -> AgentState:
    print("节点：生成图表...")
    if state['raw_data'] is None or state['predictions'] is None or state['prediction_dates'] is None:
        state['error_message'] = "数据不完整，无法生成图表。"
        state['next_action'] = 'handle_error'
        return state
    
    plot_file = plot_forecast(state['raw_data'], state['prediction_dates'], state['predictions'])
    if plot_file:
        state['plot_path'] = plot_file
        print(f"图表已生成并保存到: {state['plot_path']}")
        state['error_message'] = None
    else:
        state['error_message'] = "图表生成失败。"
    state['next_action'] = 'get_user_query' # 完成一次预测后，回到提问
    return state

def handle_error_node(state: AgentState) -> AgentState:
    print(f"发生错误: {state.get('error_message', '未知错误')}")
    # 简单处理：回到用户查询或结束
    state['next_action'] = 'get_user_query' 
    # 或者可以设计更复杂的错误恢复逻辑
    # state['next_action'] = 'end_conversation'
    return state

def end_node(state: AgentState) -> AgentState:
    print("对话结束。感谢使用！")
    return state # 最终状态

# --- 4. 构建 LangGraph 图 ---
workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("initialize", initialize_agent)
workflow.add_node("load_data", load_data_node)
workflow.add_node("train_model", train_model_node)
workflow.add_node("get_user_query", get_user_query_node)
workflow.add_node("predict_exchange_rate", predict_exchange_rate_node)
workflow.add_node("generate_plot", generate_plot_node)
workflow.add_node("handle_error", handle_error_node)
workflow.add_node("end_conversation", end_node) # 明确的结束节点

# 设置入口点
workflow.set_entry_point("initialize")

# 定义边的条件逻辑
def decide_next_action(state: AgentState) -> str:
    return state.get('next_action', 'end_conversation') # 默认结束

# 添加条件边
workflow.add_conditional_edges(
    "initialize",
    decide_next_action,
    {
        "load_data": "load_data",
        "end_conversation": END # 如果初始化失败或直接结束
    }
)
workflow.add_conditional_edges(
    "load_data",
    decide_next_action,
    {
        "train_model": "train_model",
        "handle_error": "handle_error",
        "end_conversation": END
    }
)
workflow.add_conditional_edges(
    "train_model",
    decide_next_action,
    {
        "get_user_query": "get_user_query",
        "handle_error": "handle_error",
        "end_conversation": END
    }
)
workflow.add_conditional_edges(
    "get_user_query",
    decide_next_action,
    {
        "predict_exchange_rate": "predict_exchange_rate",
        "get_user_query": "get_user_query", # 循环提问
        "end_conversation": END
    }
)
workflow.add_conditional_edges(
    "predict_exchange_rate",
    decide_next_action,
    {
        "generate_plot": "generate_plot",
        "handle_error": "handle_error",
        "end_conversation": END
    }
)
workflow.add_conditional_edges(
    "generate_plot",
    decide_next_action,
    {
        "get_user_query": "get_user_query", # 完成后回到提问
        "handle_error": "handle_error", # 如果绘图失败
        "end_conversation": END
    }
)
workflow.add_conditional_edges(
    "handle_error",
    decide_next_action, # 错误处理后决定去哪里
    {
        "get_user_query": "get_user_query",
        "load_data": "load_data", # 例如，尝试重新加载数据
        "train_model": "train_model",
        "end_conversation": END
    }
)

# 编译图
app = workflow.compile()

# --- 5. 运行 LangGraph 应用 ---
if __name__ == '__main__':
    print("启动 LangGraph 汇率预测代理...")
    # 初始状态可以为空字典，或者包含一些初始值
    initial_state = {}
    
    # LangGraph 的 invoke 通常用于单次执行，对于持续对话，需要一个循环
    # 或者使用 stream 方法并处理事件流
    # 这里我们简化为一次完整的流程，然后可以手动再次调用 app.invoke
    
    final_state = app.invoke(initial_state)
    
    # 如果需要持续交互，可以包裹在一个循环中，并根据 final_state['next_action'] 来决定是否继续
    # 例如:
    # current_state = initial_state
    # while True:
    #     current_state = app.invoke(current_state)
    #     if current_state.get('next_action') == 'end_conversation' or app.config_schema().is_end(current_state):
    #         break
    #     # 可能需要在这里重置一些状态或获取新的用户输入来驱动下一次迭代

    print("\n代理运行结束。最终状态:")
    # print(final_state) # 打印最终状态以供调试